#include "md501.h"
#include <iomanip>
#include <assert.h>
#include <chrono>
#include <arm_neon.h>

using namespace std;
using namespace chrono;

Byte **StringProcessBatch(const string inputs[4], int n_bytes[4])
{
    Byte **paddedMessages = new Byte *[4];
    for (int idx = 0; idx < 4; ++idx)
    {
        if (inputs[idx].empty())
        {
            paddedMessages[idx] = nullptr;
            n_bytes[idx] = 0;
            continue;
        }

        int length = inputs[idx].length();
        int bitLength = length * 8;
        int paddingBits = bitLength % 512;

        if (paddingBits > 448)
            paddingBits += 512 - (paddingBits - 448);
        else if (paddingBits < 448)
            paddingBits = 448 - paddingBits;
        else
            paddingBits = 512;

        int paddingBytes = paddingBits / 8;
        int paddedLength = length + paddingBytes + 8;
        paddedMessages[idx] = new Byte[paddedLength];

        memcpy(paddedMessages[idx], inputs[idx].c_str(), length);
        paddedMessages[idx][length] = 0x80;
        memset(paddedMessages[idx] + length + 1, 0, paddingBytes - 1);

        for (int i = 0; i < 8; ++i)
        {
            paddedMessages[idx][length + paddingBytes + i] = ((uint64_t)bitLength >> (i * 8)) & 0xFF;
        }

        n_bytes[idx] = paddedLength;
    }
    return paddedMessages;
}
//g++ main.cpp train.cpp guessing.cpp md502.cpp -o main -O2
void MD5Hash(const string inputs[4], uint32_t states[4][4])
{
    for(int j=0;j<4;j++){
    cout<<inputs[j]<<" ";}
    // 初始化四路状态
    uint32x4_t abcd[4] = {
        vdupq_n_u32(0x67452301), // a
        vdupq_n_u32(0xefcdab89), // b
        vdupq_n_u32(0x98badcfe), // c
        vdupq_n_u32(0x10325476)  // d
    };
    int *n_bytes=new int[1];
    cout<<"n_bytes="<<n_bytes<<" ";
    Byte **paddedMessages = StringProcessBatch(inputs, n_bytes);
    // 为四路并行创建临时变量
    cout<<"processed"<<" ";
    uint32x4_t a, b, c, d;
    Byte *message;
    int n_blocks;
    for (int idx = 0; idx < 4; ++idx)
    {
        cout<<"idx="<<idx<<" ";
        message = paddedMessages[idx];
    }
    n_blocks = n_bytes[0] / 64;
    for (int block_idx = 0; block_idx < n_blocks; ++block_idx)
    {
        cout<<"block_idx="<<block_idx<<" ";
        // 加载当前block的16个words到x数组
        bit32 x[16];
        for (int i = 0; i < 16; ++i)
        {
            x[i] = (message[4 * i + block_idx * 64]) |
                   (message[4 * i + 1 + block_idx * 64] << 8) |
                   (message[4 * i + 2 + block_idx * 64] << 16) |
                   (message[4 * i + 3 + block_idx * 64] << 24);
        }

        // 加载到NEON寄存器（四路广播）
        uint32x4_t x_simd[16];
        for (int i = 0; i < 16; ++i)
        {
            x_simd[i] = vdupq_n_u32(x[i]);
        }

        // 保存初始状态
        a = abcd[0];
        b = abcd[1];
        c = abcd[2];
        d = abcd[3];

        /* Round 1 */
        FF_SIMD(a, b, c, d, x_simd[0], 7, vdupq_n_u32(0xd76aa478));
        FF_SIMD(d, a, b, c, x_simd[1], 12, vdupq_n_u32(0xe8c7b756));
        FF_SIMD(c, d, a, b, x_simd[2], 17, vdupq_n_u32(0x242070db));
        FF_SIMD(b, c, d, a, x_simd[3], 22, vdupq_n_u32(0xc1bdceee));

        FF_SIMD(a, b, c, d, x_simd[4], 7, vdupq_n_u32(0xf57c0faf));
        FF_SIMD(d, a, b, c, x_simd[5], 12, vdupq_n_u32(0x4787c62a));
        FF_SIMD(c, d, a, b, x_simd[6], 17, vdupq_n_u32(0xa8304613));
        FF_SIMD(b, c, d, a, x_simd[7], 22, vdupq_n_u32(0xfd469501));

        FF_SIMD(a, b, c, d, x_simd[8], 7, vdupq_n_u32(0x698098d8));
        FF_SIMD(d, a, b, c, x_simd[9], 12, vdupq_n_u32(0x8b44f7af));
        FF_SIMD(c, d, a, b, x_simd[10], 17, vdupq_n_u32(0xffff5bb1));
        FF_SIMD(b, c, d, a, x_simd[11], 22, vdupq_n_u32(0x895cd7be));

        FF_SIMD(a, b, c, d, x_simd[12], 7, vdupq_n_u32(0x6b901122));
        FF_SIMD(d, a, b, c, x_simd[13], 12, vdupq_n_u32(0xfd987193));
        FF_SIMD(c, d, a, b, x_simd[14], 17, vdupq_n_u32(0xa679438e));
        FF_SIMD(b, c, d, a, x_simd[15], 22, vdupq_n_u32(0x49b40821));

        /* Round 2 */
        GG_SIMD(a, b, c, d, x_simd[1], 5, vdupq_n_u32(0xf61e2562));
        GG_SIMD(d, a, b, c, x_simd[6], 9, vdupq_n_u32(0xc040b340));
        GG_SIMD(c, d, a, b, x_simd[11], 14, vdupq_n_u32(0x265e5a51));
        GG_SIMD(b, c, d, a, x_simd[0], 20, vdupq_n_u32(0xe9b6c7aa));

        GG_SIMD(a, b, c, d, x_simd[5], 5, vdupq_n_u32(0xd62f105d));
        GG_SIMD(d, a, b, c, x_simd[10], 9, vdupq_n_u32(0x02441453));
        GG_SIMD(c, d, a, b, x_simd[15], 14, vdupq_n_u32(0xd8a1e681));
        GG_SIMD(b, c, d, a, x_simd[4], 20, vdupq_n_u32(0xe7d3fbc8));

        GG_SIMD(a, b, c, d, x_simd[9], 5, vdupq_n_u32(0x21e1cde6));
        GG_SIMD(d, a, b, c, x_simd[14], 9, vdupq_n_u32(0xc33707d6));
        GG_SIMD(c, d, a, b, x_simd[3], 14, vdupq_n_u32(0xf4d50d87));
        GG_SIMD(b, c, d, a, x_simd[8], 20, vdupq_n_u32(0x455a14ed));

        GG_SIMD(a, b, c, d, x_simd[13], 5, vdupq_n_u32(0xa9e3e905));
        GG_SIMD(d, a, b, c, x_simd[2], 9, vdupq_n_u32(0xfcefa3f8));
        GG_SIMD(c, d, a, b, x_simd[7], 14, vdupq_n_u32(0x676f02d9));
        GG_SIMD(b, c, d, a, x_simd[12], 20, vdupq_n_u32(0x8d2a4c8a));

        /* Round 3 */
        HH_SIMD(a, b, c, d, x_simd[5], 4, vdupq_n_u32(0xfffa3942));
        HH_SIMD(d, a, b, c, x_simd[8], 11, vdupq_n_u32(0x8771f681));
        HH_SIMD(c, d, a, b, x_simd[11], 16, vdupq_n_u32(0x6d9d6122));
        HH_SIMD(b, c, d, a, x_simd[14], 23, vdupq_n_u32(0xfde5380c));

        HH_SIMD(a, b, c, d, x_simd[1], 4, vdupq_n_u32(0xa4beea44));
        HH_SIMD(d, a, b, c, x_simd[4], 11, vdupq_n_u32(0x4bdecfa9));
        HH_SIMD(c, d, a, b, x_simd[7], 16, vdupq_n_u32(0xf6bb4b60));
        HH_SIMD(b, c, d, a, x_simd[10], 23, vdupq_n_u32(0xbebfbc70));

        HH_SIMD(a, b, c, d, x_simd[13], 4, vdupq_n_u32(0x289b7ec6));
        HH_SIMD(d, a, b, c, x_simd[0], 11, vdupq_n_u32(0xeaa127fa));
        HH_SIMD(c, d, a, b, x_simd[3], 16, vdupq_n_u32(0xd4ef3085));
        HH_SIMD(b, c, d, a, x_simd[6], 23, vdupq_n_u32(0x04881d05));

        HH_SIMD(a, b, c, d, x_simd[9], 4, vdupq_n_u32(0xd9d4d039));
        HH_SIMD(d, a, b, c, x_simd[12], 11, vdupq_n_u32(0xe6db99e5));
        HH_SIMD(c, d, a, b, x_simd[15], 16, vdupq_n_u32(0x1fa27cf8));
        HH_SIMD(b, c, d, a, x_simd[2], 23, vdupq_n_u32(0xc4ac5665));

        /* Round 4 */
        II_SIMD(a, b, c, d, x_simd[0], 6, vdupq_n_u32(0xf4292244));
        II_SIMD(d, a, b, c, x_simd[7], 10, vdupq_n_u32(0x432aff97));
        II_SIMD(c, d, a, b, x_simd[14], 15, vdupq_n_u32(0xab9423a7));
        II_SIMD(b, c, d, a, x_simd[5], 21, vdupq_n_u32(0xfc93a039));

        II_SIMD(a, b, c, d, x_simd[12], 6, vdupq_n_u32(0x655b59c3));
        II_SIMD(d, a, b, c, x_simd[3], 10, vdupq_n_u32(0x8f0ccc92));
        II_SIMD(c, d, a, b, x_simd[10], 15, vdupq_n_u32(0xffeff47d));
        II_SIMD(b, c, d, a, x_simd[1], 21, vdupq_n_u32(0x85845dd1));

        II_SIMD(a, b, c, d, x_simd[8], 6, vdupq_n_u32(0x6fa87e4f));
        II_SIMD(d, a, b, c, x_simd[15], 10, vdupq_n_u32(0xfe2ce6e0));
        II_SIMD(c, d, a, b, x_simd[6], 15, vdupq_n_u32(0xa3014314));
        II_SIMD(b, c, d, a, x_simd[13], 21, vdupq_n_u32(0x4e0811a1));

        II_SIMD(a, b, c, d, x_simd[4], 6, vdupq_n_u32(0xf7537e82));
        II_SIMD(d, a, b, c, x_simd[11], 10, vdupq_n_u32(0xbd3af235));
        II_SIMD(c, d, a, b, x_simd[2], 15, vdupq_n_u32(0x2ad7d2bb));
        II_SIMD(b, c, d, a, x_simd[9], 21, vdupq_n_u32(0xeb86d391));

        // 更新状态
        abcd[0] = vaddq_u32(abcd[0], a);
        abcd[1] = vaddq_u32(abcd[1], b);
        abcd[2] = vaddq_u32(abcd[2], c);
        abcd[3] = vaddq_u32(abcd[3], d);

        // 存储结果
        vst1q_u32(&states[0][0], a);
        vst1q_u32(&states[1][0], b);
        vst1q_u32(&states[2][0], c);
        vst1q_u32(&states[3][0], d);
    }
        // 字节序调整
        for (int idx = 0; idx < 4; ++idx)
        {
            for (int i = 0; i < 4; ++i)
            {
                uint32_t value = states[idx][i];
                states[idx][i] = ((value & 0xFF) << 24) |
                                 ((value & 0xFF00) << 8) |
                                 ((value & 0xFF0000) >> 8) |
                                 ((value & 0xFF000000) >> 24);
            }
        }
    
    delete[] paddedMessages;
    delete[] n_bytes;
    cout<<"hashed"<<endl<<std::flush;
}